package cn.feizhu.aimarket.ai.workflow.config.process;

import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.streaming.StreamingOutput;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.codec.ServerSentEvent;
import reactor.core.publisher.Sinks;

import java.util.Map;
import java.util.concurrent.CompletionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * 智能内容生成工作流流式处理器
 */
public class WorkflowGraphProcess {

    private static final Logger logger = LoggerFactory.getLogger(WorkflowGraphProcess.class);

    private final ExecutorService executor = Executors.newSingleThreadExecutor();
    private CompiledGraph compiledGraph;

    public WorkflowGraphProcess(CompiledGraph compiledGraph) {
        this.compiledGraph = compiledGraph;
    }

    public void processStream(AsyncGenerator<NodeOutput> generator, Sinks.Many<ServerSentEvent<String>> sink) {
        executor.submit(() -> {
            generator.forEachAsync(output -> {
                try {
                    logger.info("Workflow node output = {}", output);
                    String nodeName = output.node();
                    String content;
                    if (output instanceof StreamingOutput streamingOutput) {
                        // 处理流式输出（如AI生成的内容） 使用Map.of创建key和value都不能为null
                        if(nodeName == null) {
                            nodeName = "";
                        }
                        content = JSON.toJSONString(Map.of(
                            "node", nodeName,
                            "type", "streaming",
                            "chunk", streamingOutput.chunk() == null ? "" : streamingOutput.chunk()
                        ));
                    } else {
                        // 处理普通节点输出
                        JSONObject nodeOutput = new JSONObject();
                        nodeOutput.put("node", nodeName);
                        nodeOutput.put("type", "complete");
                        nodeOutput.put("data", output.state().data());
                        // TODO 这里根据节点名来处理

                        // 根据节点类型添加特殊处理
                        switch (nodeName) {
                            case WorkflowControlConstant.SMART_CHECK:
                                nodeOutput.put("message", "智能判断完成");
                                break;
                            case WorkflowControlConstant.COMMAND_HANDLE:
                                nodeOutput.put("message", output.state().value(WorkflowControlConstant.PROCESSED_COMMAND));
                                break;
                            case "outline_generate":
                                nodeOutput.put("message", "大纲生成完成");
                                break;
                            case WorkflowControlConstant.OUTLINE_CHECK:
                                nodeOutput.put("message", "大纲验证完成");
                                break;
                            case "outline_fix":
                                nodeOutput.put("message", "大纲修复完成");
                                break;
                            case WorkflowControlConstant.HUMAN_FEEDBACK:
                                nodeOutput.put("message", output.state().value(WorkflowControlConstant.FEEDBACK));
                                break;
                            case WorkflowControlConstant.MULTI_GENERATE:
                                nodeOutput.put("message", "多文档生成完成");
                                break;
                            case WorkflowControlConstant.AGGREGATE:
                                nodeOutput.put("message", "内容聚合完成");
                                break;
                            case WorkflowControlConstant.VISUALIZE:
                                nodeOutput.put("message", "可视化生成完成");
                                break;
                            case WorkflowControlConstant.RETURN_URL:
                                nodeOutput.put("message", output.state().value(WorkflowControlConstant.VISUAL_URL));
                                break;
                            default:
                                nodeOutput.put("message", "节点处理完成");
                        }

                        content = JSON.toJSONString(nodeOutput);
                    }

                    sink.tryEmitNext(ServerSentEvent.builder(content).build());
                } catch (Exception e) {
                    logger.error("处理节点输出时发生错误", e);
                    throw new CompletionException(e);
                }
            }).thenAccept(v -> {
                // TODO 暂停的时候也会触发，所以不能说是完成，只有返回url才可以通知前端完成
                logger.info("工作流流式处理完成");
                sink.tryEmitNext(ServerSentEvent.builder(
                    JSON.toJSONString(Map.of("type", "complete"))
                ).build());
                sink.tryEmitComplete();
            }).exceptionally(e -> {
                logger.error("工作流流式处理异常", e);
//                sink.tryEmitError(e);
                return null;
            });
        });
    }
}
